from torch import nn
from torch_geometric.nn import GraphConv


# 定义一个时空模块（Spatial-Temporal Block），用于处理时空数据
class STBlock(nn.Module):
    def __init__(self, in_channels, out_channels, A):
        super().__init__()
        # 时序卷积层，用于处理时间维度的信息
        # 输入通道数为in_channels，输出通道数为out_channels
        # 卷积核大小为(3,1)，在时间维度上卷积，空间维度上不卷积
        # padding=(1,0)用于保持时间维度的长度不变
        self.temporal = nn.Conv2d(in_channels, out_channels, (3, 1), padding=(1, 0))

        # 图卷积层，用于处理空间维度的信息
        # 输入通道数和输出通道数均为out_channels
        # A是邻接矩阵，表示节点之间的连接关系
        self.graph = GraphConv(out_channels, out_channels, A)

    def forward(self, x):
        # 先进行时序卷积，处理时间维度的信息
        x = self.temporal(x)  # 输入形状为 [B, C, T, N]，B是批次大小，C是通道数，T是时间步长，N是节点数
        # 再进行图卷积，处理空间维度的信息
        x = self.graph(x)  # 输出形状仍为 [B, C, T, N]
        return x


# 定义一个时空图神经网络（Spatial-Temporal Graph Neural Network）
class STGNN(nn.Module):
    def __init__(self, A, num_classes=None):
        super().__init__()
        # 邻接矩阵A，表示节点之间的连接关系
        self.A = A
        # 第一个时空模块，输入通道数为16，输出通道数为32
        self.block1 = STBlock(16, 32, A)
        # 第二个时空模块，输入通道数为32，输出通道数为64
        self.block2 = STBlock(32, 64, A)
        # 全连接层，用于将特征映射到类别空间
        # 注意：这里需要定义num_classes，表示类别数
        self.fc = nn.Linear(64, num_classes)

    def forward(self, x):
        # 通过第一个时空模块
        x = self.block1(x)
        # 通过第二个时空模块
        x = self.block2(x)
        # 对时间维度进行平均，得到每个节点在所有时间步长上的平均特征
        x = x.mean(dim=2)  # 输入形状为 [B, C, T, N]，输出形状为 [B, C, N]
        # 调整维度顺序，准备输入全连接层
        x = self.fc(x.permute(0, 2, 1))  # [B, C, N] → [B, N, C] → [B, N, Class]
        return x